import os
import time
import threading
import multiprocessing

import redis


# redis所在主机的ip,默认127.0.0.1
HOST = "127.0.0.1"

# redis监听的端口,默认6379
PORT = 6379

# 需要连接的数据库,默认0号数据库
DB = 0

# 连接redis的密码,按实际配置,默认无密码
PWD = None

# 添加黑名单和解封禁黑名单以互相独立的两个任务进行处理,建议使用多进程
# 使用多线程: thread; 使用多进程: process
MULTI_MODE = "process"

# 保存需要短时间封禁的ip黑名单的ipset名称,按实际配置;
# 如果没有在服务器提前创建好该名称的ipset,程序将以hash:net模式,ipset默认的最大ip名单容量,自动创建该ipset
IPSET_NAME = "blacklists"

# 保存需要永久性封禁的ip黑名单的ipset名称,如果不进行配置,默认与IPSET_NAME配置项保持一致
# 需要永久性封禁的ip,在往redis中添加数据数据时,需要为键(key)添加后缀,添加方式: "-" + 后缀名称;
BLACK_IP_FOREVER_IPSET_NAME =  "blacklists-forever"

# 禁止黑名单访问的端口,默认80端口
LIMIT_PORT = 80
# True: 程序将创建限制 ipset 中的ip名单进行访问 端口{LIMIT_PORT} 的 iptables 规则
# False: 程序不进行任何访问规则的创建.需要开发者提前手动在服务器进行配置,此时 LIMIT_PORT 不生效
MAKE_ACCESS_RUlE = True


def make_conn():
    """
    建立redis连接
    """
    # 因为该脚本的启动需要连接redis, 在使用进程管理工具同时管理该程序及redis时,redis可能会更晚完成启动
    # 捕捉异常,如果连接redis失败,休眠一小段时间,再递归执行,以此保证程序的正常运行
    params = dict()
    params["host"] = HOST if HOST else "127.0.0.1"
    params["port"] = PORT if PORT else 6379
    params["db"] = DB if DB else 0
    if PWD:
        params["password"] = PWD
    try:
        conn = redis.StrictRedis(**params)
    except Exception:
        time.sleep(0.5)
        conn = make_conn()
    return conn, params["db"]


def add_blacklist(forever_ipset):
    """
    订阅redis添加string(k-v)事件,添加黑名单
    """
    conn, db = make_conn()
    pub = conn.pubsub()
    pub.subscribe('__keyevent@{}__:set'.format(db))
    for msg in pub.listen():
        black_key = msg["data"]
        # 刚启动程序时,如果订阅成功,会接收到一个响应消息, data字段为 1, 整型
        if type(black_key) != int:
            black_key = black_key.decode()
            ip = black_key.split("-")
            if len(ip) == 1:
                ipset_name = IPSET_NAME
            else:
                ipset_name = forever_ipset
                conn.delete(black_key)
            os.popen("sudo ipset add {} {}".format(ipset_name, ip[0]))


def del_blacklist():
    """
    订阅redis过期事件,删除黑名单
    """
    conn, db = make_conn()
    pub = conn.pubsub()
    pub.subscribe('__keyevent@{}__:expired'.format(db))
    for msg in pub.listen():
        black_key = msg["data"]
        if type(black_key) != int:
            ip = black_key.decode()
            os.popen("sudo ipset del {} {}".format(IPSET_NAME, ip))


if __name__ == '__main__':

    if not IPSET_NAME:
        raise AttributeError("A missing name of ipset is necessary")

    params = [BLACK_IP_FOREVER_IPSET_NAME if BLACK_IP_FOREVER_IPSET_NAME else IPSET_NAME,]
    add_dict = dict(target=add_blacklist, args=params)

    if MULTI_MODE == "process":
        add_task = multiprocessing.Process(**add_dict)
        del_task = multiprocessing.Process(target=del_blacklist,)

    elif MULTI_MODE == "thread":
        add_task = threading.Thread(**add_dict)
        del_task = threading.Thread(target=del_blacklist,)

    else:
        raise AttributeError("The MULTI_MODE must be choosed in 'process' or 'thread'")

    limit_port = LIMIT_PORT if LIMIT_PORT else 80
    for ipset_name in [IPSET_NAME, BLACK_IP_FOREVER_IPSET_NAME]:
        # 如果ipset已经创建完成,再次创建操作将会失败,程序不会停止执行
        if ipset_name:
            os.popen("sudo ipset create {} hash:net".format(ipset_name))
            if MAKE_ACCESS_RUlE == True:
                os.popen("sudo iptables -I INPUT -m set --match-set  {}  src -p tcp --destination-port {} -j DROP".format(ipset_name, limit_port))

    add_task.start()
    del_task.start()
